import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets

#x: [60k,28,28]
#y:[60k]
(x,y),(x_test,y_test) = datasets.mnist.load_data()

x = tf.convert_to_tensor(x, dtype=tf.float32) / 255.
x_test = tf.convert_to_tensor(x_test, dtype=tf.float32) / 255.
y = tf.convert_to_tensor(y, dtype=tf.int32)
y_test = tf.convert_to_tensor(y_test, dtype=tf.int32)

print(x.shape,y.shape)
print(tf.reduce_min(x),tf.reduce_max(x))
print(tf.reduce_min(y),tf.reduce_max(y))

#创建数据集
train_db = tf.data.Dataset.from_tensor_slices((x,y)).batch(128)
train_iter = iter(train_db)#迭代器
test_db = tf.data.Dataset.from_tensor_slices((x_test,y_test)).batch(128)
sample = next(train_iter)

print('batch:',sample[0].shape)

#[b,781] => [b,512] => [b,128] => [b,10]
# 如果参数不合适。会出现梯度爆炸的现象，所以调参方差0.1
# 计算梯度。需要参数为Variable类型
w1 = tf.Variable(tf.random.truncated_normal([784,256], stddev=0.1))
b1 = tf.Variable(tf.zeros([256]))
w2 = tf.Variable(tf.random.truncated_normal([256,128],stddev=0.1))
b2 = tf.Variable(tf.zeros([128]))
w3 = tf.Variable(tf.random.truncated_normal([128,10],stddev=0.1))
b3 = tf.Variable(tf.zeros([10]))

learning_rate = 1e-3
for epoch in range(100):#对整个数据集迭代10次
    for step,(x,y) in enumerate(train_db):
        #x [ 128,28,28]
        #y [128]
        x = tf.reshape(x,[-1,28*28])

        with tf.GradientTape() as tape:
            #x [b,28*28]
            # h1 = x@w1 + b1
            h1 = x@w1 + b1
            h1 = tf.nn.relu(h1)
            h2 = h1@w2 + b2
            h2 = tf.nn.relu(h2)
            out = h2@w3 + b3

            # 计算误差
            # out[b,10]
            # y : [b] => [b,10]
            y_onehot = tf.one_hot(y,depth=10)

            # mse = mean((y-out)^2）
            #[b,10]
            loss = tf.square(y_onehot - out)
            # mean:scalar
            loss = tf.reduce_mean(loss)

        # 梯度计算
        grads = tape.gradient(loss, [w1,b1,w2,b2,w3,b3])
        #w1 = w1 - lr * w1_grad
        w1.assign_sub(learning_rate * grads[0])  #原地跟新
        b1.assign_sub(learning_rate * grads[1])  # 原地跟新
        w2.assign_sub(learning_rate * grads[2])  # 原地跟新
        b2.assign_sub(learning_rate * grads[3])  # 原地跟新
        w3.assign_sub(learning_rate * grads[4])  # 原地跟新
        b3.assign_sub(learning_rate * grads[5])  # 原地跟新
        # w1 = w1 - learning_rate * grads[0]
        # b1 = b1 - learning_rate * grads[1]
        # w2 = w2 - learning_rate * grads[2]
        # b2 = b2 - learning_rate * grads[3]
        # w3 = w3 - learning_rate * grads[4]
        # b3 = b3 - learning_rate * grads[5]

        if step % 100 == 0:
            #打印损失
            print('epoch:{0} step{1}: loss={2}'.format(epoch, step,loss))

    total_correct = 0
    total_num = 0
    for x,y in test_db:

        # [b,28,28]
        x = tf.reshape(x,[-1,28*28])

        #[b,784] => [b,256] => [b,128] => [b,10]
        h1 = tf.nn.relu(x@w1 + b1)
        h2 = tf.nn.relu(h1@w2 + b2)
        out = h2@w3 + b3

        # out : [b,10] ~ R
        # prob [b,10] ~[0,1]
        prob = tf.nn.softmax(out,axis=1)
        #[b,10] => [b]
        pred = tf.argmax(prob,axis=1)
        pred = tf.cast(pred,dtype=tf.int32)
        # y: [b]
        correct = tf.cast(tf.equal(pred,y),dtype=tf.int32)
        correct = tf.reduce_sum(correct)

        total_correct += int(correct)
        total_num += x.shape[0]


    acc = total_correct / total_num
    print('准确度为:{0}'.format(acc))